import matplotlib.pyplot as plt
import torch
import torchvision
from models.autoencoder import Autoencoder

# 加载测试数据
test_data = torchvision.datasets.CIFAR10(
    root='./data',
    train=False,
    download=True,
    transform=torchvision.transforms.ToTensor()
)

# 加载模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Autoencoder().to(device)
model.load_state_dict(torch.load("autoencoder.pth", map_location=device, weights_only=True))
model.eval()

# 可视化对比
with torch.no_grad():
    test_img = test_data[0][0].unsqueeze(0).to(device)
    reconstructed = model(test_img).cpu()

plt.figure(figsize=(8, 4))
plt.subplot(1, 2, 1)
plt.imshow(test_data[0][0].permute(1, 2, 0))
plt.title("Original")
plt.axis('off')

plt.subplot(1, 2, 2)
plt.imshow(reconstructed.squeeze().permute(1, 2, 0))
plt.title("Reconstructed")
plt.axis('off')
plt.show()
